package com.huiztech.construction.mcp.advisor;

import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import reactor.core.publisher.Flux;

import java.util.Map;


/**
 * 自定义拦截器
 */
public class ReReadingAdvisor implements BaseAdvisor {

    private static final String DEFAULT_USER_TEXT_ADVISE= """
            {re2_input_query}
            Read the question again: {re2_input_query}
            """;


//    @Override
//    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
//        // 1. 前置增强
//        advisedRequest = this.before(advisedRequest);
//
//        // 2. 执行后续的advisor链
//        Flux<AdvisedResponse> advisedResponseFlux = chain.nextAroundStream(advisedRequest);
//        // 3. 后置增强
//        this.observeAfter(advisedResponseFlux);
//        return advisedResponseFlux;
//    }
//
//
//    /**
//     * 前置增强，记录一个初始时间，并将时间存放到advisorContext
//     *
//     * @param advisedRequest
//     * @return
//     */
//    private AdvisedRequest before(AdvisedRequest advisedRequest){
//        System.out.println("前置 AdvisorRequest 增强");
//        // 存储耗时，上文中我们提到AdvisorContext会共享状态。我们将耗时记录在上下文中，当然也可以自己定义一个变量存储
//        advisedRequest.adviseContext().put("start0", System.currentTimeMillis());
//        return advisedRequest;
//    }
//
//    /**
//     * 后置增强，获取初始时间，并计算耗时
//     * @param advisedResponses
//     */
//    private void observeAfter(Flux<AdvisedResponse> advisedResponses){
//        System.out.println("后置 advisedResponse 增强");
//        advisedResponses.toIterable().forEach(advisedResponse -> {
//            Long start0 = (Long) advisedResponse.adviseContext().get("start0");
//        System.out.println("从 AdvisorContext中获取start0:" + start0);
//        System.out.println("耗时：" + (System.currentTimeMillis() - start0));
//        });
//    }


//    @Override
//    public String getName() {
//        return "RereadingAdvisor";
//    }

    /**
     * 强调用户输入的提示词
     * @param chatClientRequest
     * @param advisorChain
     * @return
     */
    @Override
    public ChatClientRequest before(ChatClientRequest chatClientRequest, AdvisorChain advisorChain) {
        // 用户输入文本
        String text = chatClientRequest.prompt().getUserMessage().getText();
        // 定义重复输入模板
        String augmentedSystemText = PromptTemplate.builder().template(DEFAULT_USER_TEXT_ADVISE).build().render(Map.of("re2_input_query", text));

        //设置请求的提示词
        ChatClientRequest processedChatClientRequest =
                // 不保留
                ChatClientRequest.builder()
                        .prompt(Prompt.builder().content(augmentedSystemText).build())
                        .build();
        return processedChatClientRequest;
    }

    @Override
    public ChatClientResponse after(ChatClientResponse chatClientResponse, AdvisorChain advisorChain) {

        return chatClientResponse;
    }

    @Override
    public int getOrder() {
        return 0;
    }

}
